# Copyright (c) OpenMMLab. All rights reserved.
from typing import List
import os.path as osp
import logging
import mmengine
from mmengine.logging import print_log
import mmengine.fileio as fileio

from mmseg.registry import DATASETS
from .basesegdataset import BaseSegDataset


@DATASETS.register_module()
class MainVesselDataset(BaseSegDataset):
    """冠状动脉双分支分割数据集 (LCX + LAD).

    数据目录示例:
    dataset/
    ├── images/
    │   ├── training/
    │   └── validation/
    └── labels/
        ├── training/
        │   ├── lcx/
        │   └── lad/
        └── validation/
            ├── lcx/
            └── lad/

    每个样本会返回一张图像和两张 mask (lcx, lad)，用于多头分割模型。
    """

    METAINFO = dict(
        classes=('background', 'vessel'),
        palette=[[120, 120, 120], [6, 230, 230]]
    )

    def __init__(self,
                 img_suffix='.png',
                 seg_map_suffix='.png',
                 reduce_zero_label=False,
                 data_prefix=None,
                 **kwargs) -> None:
        # 解析 data_prefix
        img_dir = data_prefix.get('img_path')
        ann_dirs = data_prefix.get('seg_map_path', None)

        assert isinstance(ann_dirs, (list, tuple)) and len(ann_dirs) == 2, \
            '`seg_map_path` 必须是 [lcx_dir, lad_dir]'

        self.lcx_dir, self.lad_dir = ann_dirs[0], ann_dirs[1]

        # 只把 img_dir 传给 BaseSegDataset，避免它尝试处理 list
        super().__init__(
            img_suffix=img_suffix,
            seg_map_suffix=seg_map_suffix,
            reduce_zero_label=reduce_zero_label,
            data_prefix=dict(img_path=img_dir),
            **kwargs)

        # assert fileio.exists(img_dir, backend_args=self.backend_args)

    def load_data_list(self) -> List[dict]:
        """加载图像及 LCX/LAD 标签路径."""
        data_list = []
        img_dir = self.data_prefix.get('img_path', None)
        _suffix_len = len(self.img_suffix)

        for img in fileio.list_dir_or_file(
                dir_path=img_dir,
                list_dir=False,
                suffix=self.img_suffix,
                recursive=True,
                backend_args=self.backend_args):
            data_info = dict(img_path=osp.join(img_dir, img))
            seg_map_name = img[:-_suffix_len] + self.seg_map_suffix

            # LCX 和 LAD 标签分别保存
            data_info['seg_map_path_lcx'] = osp.join(self.data_root, self.lcx_dir, seg_map_name)
            data_info['seg_map_path_lad'] = osp.join(self.data_root, self.lad_dir, seg_map_name)

            data_info['label_map'] = self.label_map
            data_info['reduce_zero_label'] = self.reduce_zero_label
            data_info['seg_fields'] = []
            data_list.append(data_info)

        data_list = sorted(data_list, key=lambda x: x['img_path'])
        return data_list

    def get_gt_seg_maps_by_idx(self, index: int) -> dict:
        """返回 index 样本的 LCX 和 LAD mask."""
        data_info = self.data_list[index]
        results = {}
        if 'seg_map_path_lcx' in data_info:
            results['gt_seg_map_lcx'] = self.file_client.get(
                data_info['seg_map_path_lcx'], backend_args=self.backend_args)
        if 'seg_map_path_lad' in data_info:
            results['gt_seg_map_lad'] = self.file_client.get(
                data_info['seg_map_path_lad'], backend_args=self.backend_args)
        return results
